from torch.nn import init


### initalize the module
def init_weights(net, init_type="normal"):
    # print('initialization method [%s]' % init_type)
    if init_type == "kaiming":
        net.apply(weights_init_kaiming)
    else:
        raise NotImplementedError(
            "initialization method [%s] is not implemented" % init_type
        )


def weights_init_kaiming(m):
    classname = m.__class__.__name__
    # print(classname)
    if classname.find("Conv") != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
    elif classname.find("Linear") != -1:
        init.kaiming_normal_(m.weight.data, a=0, mode="fan_in")
    elif classname.find("BatchNorm") != -1:
        init.normal_(m.weight.data, 1.0, 0.02)
        init.constant_(m.bias.data, 0.0)


### compute model params
def count_param(model):
    return sum(param.view(-1).size()[0] for param in model.parameters())
